import numpy as np
import torchvision.datasets as datasets
import torch
import torch.nn as nn
from torchvision import transforms
from torch.utils.data.dataloader import DataLoader
import torch.optim as optim
import matplotlib.pylab as plt
import tqdm
# from timm.utils import accuracy, AverageMeter
Cuda=torch.cuda.is_available()
plt.switch_backend('agg')
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

# 设置transforms
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),  # numpy -> Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化 ，范围[-1,1]
])
# 训练集
trainset = datasets.CIFAR10(root='data/CIFAR10', train=True, download=False, transform=transform)
# 测试集

BATCH_SIZE = 8

train_loader = DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
# test_loader = DataLoader(testset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
from model.Swim_transformer import *

model=Swin_Model()

model = model.train()
if Cuda:
    model = model.cuda()
optimizer = optim.AdamW(model.parameters(), lr=0.001)
Loss = nn.CrossEntropyLoss()
list1 = []
list2 = []

len_data=len(trainset)
print(len_data)
print(torch.cuda.is_available()) # 查看CUDA是否可用
if __name__=='__main__':
    histloss=[]
    histacc=[]
    for epoch in range(5):
        index=0
        acc=0.0
        accsum=0.0
        for (x,y) in tqdm.tqdm(train_loader):
            if Cuda:
                x=x.cuda()
                y=y.cuda()
            index+=1
            y_pred=model(x)
            optimizer.zero_grad()
            loss1=Loss(y_pred,y)
            yy=torch.argmax(y_pred,dim=-1)
            acc+=sum(yy==y)
            # acc1, acc5 = accuracy(y_pred, y, topk=(1, 5))
            # accsum+=acc1
            # print(acc1,acc5)
            loss1.backward()
            optimizer.step()
        acc=acc/len_data
        # accsum=accsum/index
        print(epoch,acc)